爱因斯坦校友提出的Transformer简化方案是条歧路

最近有两篇关于Transformer的论文在国内大模型圈争相转阅：

一篇是伯克利马毅教授团队白盒Transformer，针对这篇上周末笔者已经做了深度的解读 “如何看待伯克利马毅教授团队白盒Transformer”，感谢广大热心读者的支持，这篇解读创造了笔者写小作文的历史。能够有这个思想层面的系统分析，个人觉得也得益于过去半年整理的“大模型的数学物理原理的认知框架” 。笔者想借助这个框架，用本篇作文聚焦分析下面第二篇流行的论文。

这篇论文就是苏黎世联邦理工学院计算机科学系的Bobby He和Thomas Hofmann的论文“简化Transformer模块 Simplifying Transformer Blocks”。国内一些科技媒体评价是：“这是自Transformer 开始以来，看到的最好的改进”，不知道是出于真心，还是对苏黎世联邦理工学院校友爱因斯坦、冯·诺依曼、伦琴、泡利还有潘周聃的敬畏。

我们先一起回顾一下这篇论文的重要结论：

“论文以信号传播理论及实证研究结果为基础，探讨标准Transformer 块的简化方法。证明了许多组件，如跳过连接（Skip Connection）、投影或值参数(Projection or Value Paramerters)、顺序子块和归一化层（LN: Layer of Normalisation），可以在不牺牲训练速度的情况下被删除。在纯自回归解码器和纯BERT编码器模型上的实验表明，论文的简化Transformer 实现了与标准Transformer 相当的训练速度和性能，同时训练吞吐量提高了15%，使用的参数减少了15%”。

笔者天资愚顽，虽也敬畏权威，但从不迷信。通过应用自己的框架分析，笔者觉得这篇论文的简化方法是一条歧路，甚至整个简化方向都是值得商榷的，当然这个判断也包括各种量化剪枝技术。

论文作者对 Transformer的抽象符号化表达如下公式(1)-(3)：Transformer 接受输入序列Xin（d 维 T 个token），Pre-LN前置Layer Norm 输出Xout。请大家忽略其中令人生畏的复杂数学符号 FF:feed forward, SA:self attention, LN:layer normalisation, MHA:multi head attention，只看它的主干：一个线性化部分 + 一个多层感知MLP的非线性化残差部分。单Transformer 模块如此，多层堆叠在一起，就构成了嵌套迭代的套娃神经网络。

这与马毅教授团队概括的“对于原始信息x，训练过程通过编码器有损压缩成低维表征z；编码器的逆映射（解码器或生成器）从z恢复出x的近似表征x_cap”思想也一致。这跟笔者的数理认知框架“重整化提取出范畴，持续重整化驱动范畴相变，然后采样做变分推理”中的重整化部分吻合。

从“站在香农与玻尔兹曼肩上，看深度学习的术与道”中，我们知道“深度学习领域的三类最典型问题，无监督学习，有监督学习的分类与预测，归根结底都是用神经网络来近似概率分布。”所以，本质上学者们采用的这些不同表征，都是描述深度神经网络从输入信息的概率分布中，逐层提取隐变量概率分布的过程，这一过程同时也是不断从高维隐变量概率空间到低维隐变量概率空间的降维过程，伴随着变量空间尺度的不断放大，从微观到宏观。

“宇宙中生成的本质”中，笔者总结“对能量分布的辨识产生信息，无论能量在高自由度或者高维度的微观分布，还是低自由度或者低维度的宏观分布；均以分布的数学期望表达信息。对能量分布的辨识即为认知，获取能量分布的局部认知，进一步推断出原始整体分布，就是学习和智能”。香农将信息熵定义为概率分布对数的数学期望，而信息是事物现象及其属性标识的集合。因而可以说深度神经网络是事物中蕴涵信息的提取器，而提取到的是自由能概率分布。

笔者手绘这一过程示意图如下：

仔细观察这张图，深度思考这个能量概率分布的提取过程，大家可以得到三点有趣的认知：

一，现实世界中的信息概率分布是连续的（如上图中连续曲线），而深度神经网络将处理离散化为向量概率（上图中的横条），进一步将概率向量分解为，线性化部分（矩阵），和非线性化残差部分（MLP）, 借助MLP的残差部分弥补因离散化加线性化，导致的对连续概率分布拟合时候带来的非线性损失；

二，现实世界中事物的信息层次之间其实不是阶跃似的，隐变量空间尺度的不断放大也是连续的，因而神经网络Layer层之间的训练提取步长step，也就是scale 尺度的选取也应该是连续的，scale 离散化和线性化矩阵表示，同样会导致对scale 连续概率分布的拟合带来非线性部分的损失；

三，横条的多少，其实代表着隐变量空间的维度，也代表着该层对信息丰富程度的表达能力，高维可以表示低维，低维却不能表示高维。本科学习《线性代数》时候大家都学过，线性空间，可以用矩阵的线性无关列向量的维度张成，这也对应着矩阵的秩。如果矩阵的秩变成1或0, 就意味着该线性空间塌缩成一条线或点，表达能力的消失。

有了这三点认知，我们再回过去看“简化Transformer模块”这篇论文的Discussion讨论部分: "我们的深度缩放实验已经在这方面显示了希望。在理论方面，尽管去除了归一化层后，我们能够匹配 Pre-LN 块的训练速度，但关于归一化对训练速度和稳定性的益处仍然有需要回答的问题，去除归一化后我们无法获得良好的下游任务性能"。

依据上文中第一点认知，归一化其实是对笔者手绘图中的犬牙差互的横条概率之和，与平滑的概率曲线积分之间差异的校正，避免自由能概率分布偏离积分为 1。一旦取消归一化层，这个离散概率之和与连续曲线概率积分之间的偏差gap，就会在神经网络各层之间被传播和放大，从而导致训练不稳定。同时，对于这一gap，取消归一化，同样会影响MLP模块的非线性残差部分的稳定，进一步导致梯度爆炸或消失 gradient exploding or vanishing。所以“去除归一化后我们无法获得良好的下游任务性能”也就不难理解。

同样是出自苏黎世联邦理工学院的另一篇论文“Transformer中的信号传播:理论视角和秩塌缩的作用 Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapse”，还有出自其姊妹院校洛桑联邦理工学院和谷歌的另一篇论文：“注意力并不是你所需要的全部: 纯粹的注意力随着深度的增加而成倍减少  Attention is not all you need: pure attention loses rank doubly exponentially with depth ”中，讨论了"秩塌缩"问题：深度L, 宽度H 的自注意力网络，无"跳过连接"(Skip Connection)  情况下，残差矩阵的秩会以指数速率收敛到秩1，即隐变量空间塌缩成一条线。

结合上文中第二点认知，神经网络Layer层之间的训练提取step的scale尺度的选取，因离散化和线性化矩阵表示，同样会导致对scale尺度连续概率分布的拟合带来非线性部分的损失。跳过连接(Skip Connection) 通过在不同的path上跳过一些层，本质上等同于提供了不同尺度的信息提取路径，来弥补尺度选取离散化与线性化带来的非线性部分的损失。    

“简化Transformer模块”这篇论文试图从各个模块中取消跳过连接(Skip Connection) ，就会导致隐变量空间维度的获取变得不稳定，也就意味着隐变量空间的表达能力的不确定或者塌缩。根据上文第三点认知，这会表现为训练中矩阵的秩的塌缩。上文引述论文中已有论证。所以取消Skip Connection 此路不通。

我们北大学者在“重新思考Transformer和ResNets中层归一化的跳过连接(Skip Connection) Rethinking Skip Connection with Layer Normalization in Transformers and ResNets " 一文中其实总结过：“在这项工作中，我们研究了尺度因素如何影响跳过连接Skip Connection的有效性，并揭示了尺度的微小调整将导致伪梯度随着模型的深度而爆炸或消失，这可以通过归一化来解决，特别是层归一化，这导致了对普通跳过连接的一致改进”。笔者认为还是咱北大学者中肯靠谱，这一点上完胜爱因斯坦校友。

上文提到的第三点认知，其实还有一个引申意义：参数的多少，参数的精度，隐变量空间维度的大小都代表着模型提取信息的精准度，也就是对原始连续概率分布的拟合的逼近能力。增加模型的精度，增加隐变量的数量，可以扩大隐变量空间的维度，提高概率向量对实际信息的表征能力的丰富性，强化信息细微差别的区分能力。籍此笔者判断，各种DNN深度神经网络，本质上只要沿着这个思路增加参数捕获能力，都可以与Transformer殊途同归。

结合笔者的“大模型的数学物理原理的认知框架”：“重整化提取出范畴，持续重整化驱动范畴相变，然后跨范畴采样做变分推理”，这些更丰富、更精细、更大规模的参数，将能够更贴切的表达原始事物的信息。从而在其提取后的概率向量空间中，能够构建更丰富准确的（逼近连续的）范畴空间，从而为后续的可能能力涌现，以及跨范畴采样做变分推理提供更扎实的基础，一如读书多记忆好的同学写作文很能旁征博引。

所以，笔者认为，“简化Transformer模块”去除"投影或值参数(Projection or Value Paramerters)"的行为，是小模型思维，容易导致大模型褪化成小模型。各种量化剪枝降低精度的行为都适用这一判断。降参数降精度是一条歧路。

本质上，为实现对连续的概率分布的学习和分析，类脑模拟芯片或量子计算机才能根本上消除这些非线性损失导致的问题。但超出本文的议题，笔者也是一知半解，跟大家一起期待这些领域的突破。

1.SIMPLIFYING TRANSFORMER BLOCKS    https://arxiv.org/pdf/2311.01906.pdf

2.Attention is not all you need: pure attention loses rank doubly exponentially with depth  https://arxiv.org/pdf/2103.03404.pdf

3.Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapse https://arxiv.org/pdf/2206.03126.pdf

4.Rethinking Skip Connection with Layer Normalization in Transformers and ResNets    https://arxiv.org/pdf/2105.07205.pdf